using Crux.Core.LangChain.Abstractions.DocumentLoaders;
using Crux.Core.LangChain.Abstractions.Splitters.Text;

namespace Crux.Core.LangChain.Abstractions.Extensions;
/// <summary>
/// 
/// </summary>
public static class TextSplitterExtensions
{
    /// <summary>
    /// 
    /// </summary>
    /// <param name="splitter"></param>
    /// <param name="texts"></param>
    /// <param name="metadatas"></param>
    /// <returns></returns>
    /// <exception cref="ArgumentNullException"></exception>
    /// <exception cref="ArgumentException"></exception>
    public static IReadOnlyList<Document> CreateDocuments(
        this ITextSplitter splitter,
        IReadOnlyCollection<string> texts,
        IReadOnlyCollection<IDictionary<string, object>>? metadatas = null)
    {
        splitter = splitter ?? throw new ArgumentNullException(nameof(splitter));
        texts = texts ?? throw new ArgumentNullException(nameof(texts));

        var documents = new List<Document>();

        metadatas ??= Enumerable.Repeat<IDictionary<string, object>>(new Dictionary<string, object>(), texts.Count).ToList();

        if (texts.Count != metadatas.Count)
        {
            throw new ArgumentException("Number of texts and metadata must be equal.");
        }

        for (int i = 0; i < texts.Count; i++)
        {
            var text = texts.ElementAt(i);
            var metadata = metadatas.ElementAt(i);

            foreach (var chunk in splitter.SplitText(text))
            {
                documents.Add(new Document(chunk, metadata));
            }
        }

        return documents;
    }

    /// <summary>
    /// 
    /// </summary>
    /// <param name="splitter"></param>
    /// <param name="documents"></param>
    /// <returns></returns>
    public static IReadOnlyList<Document> SplitDocuments(
        this ITextSplitter splitter,
        IReadOnlyCollection<Document> documents)
    {
        var texts = documents.Select(doc => doc.PageContent).ToList();
        var metadatas = documents.Select(doc => doc.Metadata).ToList();

        return splitter.CreateDocuments(texts, metadatas);
    }


}